import json
import os
import mne
from typing import List
from copy import deepcopy

from src.tags import CustomFileToTags
from util import search_file

EVENTNAME = ['尖慢波']

def jaccard_similarity(set1:set, set2:set):
    return len(set1.intersection(set2))/len(set1.union(set2))

def get_event_duration(path, event_name=EVENTNAME, SRR=1):
    event_duration = []
    with open(path, 'r') as f:
        event_data = json.load(f)['EventData']
    for event in event_data:
        if event['EventName'] in event_name:  # in [event_name, no_mark]
            event_duration.append(set(range(int(event['Start']/SRR), int(event['Duration']/SRR))))
    return event_duration

def algorithm_merge_window(source:List[set], product:List[set], threshold=0, merge=False):
    # 重叠率/重叠区域阈值0%
    global count
    _product = deepcopy(product)
    # 窗口重叠修正
    merge_product = []
    for i in range(len(_product)-1):
        if not _product[i].isdisjoint(_product[i+1]):
            _product[i+1] = _product[i].union(_product[i+1])
        else:
            merge_product.append(_product[i])
    if len(_product) > 1:
        merge_product.append(_product[i+1])

    _true, _ptrue = 0, 0
    selected_product = merge_product if merge else product
    for s in source:
        similarity = []
        for sp in selected_product:
            similarity.append(jaccard_similarity(s, sp))
        if sum(similarity) > threshold:
            _true += 1
            _ptrue += len([s for s in similarity if s>threshold])
            # print(f'{sum(similarity)*100:2.0f}%')

    print('找到: ', _true)
    print('未找到: ', len(source)-_true)
    print('召回: %2.0f%%' %(_true/len(source)*100))
    print('精确: %2.0f%%' %(_ptrue/len(selected_product)*100))
    print('实际有: ', len(source))
    print('预测有: ', len(selected_product))
    print('-'*20)
    print(_ptrue/len(selected_product))

def get_SRR(edf_path, new_freq=100):
    raw = mne.io.read_raw_edf(edf_path)
    SRR = raw.info['sfreq']/new_freq
    raw.close()
    return SRR

def compare_tags(source_tags, product_tags, SRR=1, **kwargs):
    source = get_event_duration(source_tags, SRR=SRR)
    product = get_event_duration(product_tags, SRR=SRR)
    algorithm_merge_window(source, product, **kwargs)

def main(source_root:str, source_name:str, txt_path:str, product_root:str='./', 
         window:int=100, label_map:dict[int, str]={1:'尖慢波'}):
    out_name = os.path.basename(txt_path).split('.')[0]
    source_edf = f'{source_root}/{source_name}.edf'
    product_tags = f'{product_root}/{out_name}.tags'
    source_tags = f'{source_root}/Event/{source_name}.tags'

    SRR = get_SRR(source_edf)
    ftt = CustomFileToTags(window, label_map)
    ftt.set_SRR(SRR)
    label = ftt.load_txt(txt_path)
    ftt.auto_to_tags(label, product_root, out_name, converted=False)    
    compare_tags(source_tags, product_tags, merge=False)

if __name__ == '__main__':
    for file in search_file('../data/temp/vir/Event', exts=['tagsbak']):
        if 'v1_0' in file:
            compare_tags(file, file[:-3], 2.56)
            exit()
